import torch
from torch.utils.data import Dataset, DataLoader
import numpy as np
import ijson
from tqdm import tqdm


class flowDataset(Dataset):
    def __init__(self, dataset_json, truncate_num=100, offset=1460):
        self.trainset = []
        with open(dataset_json, 'r') as f:
            print("loading json", dataset_json)
            for json_flow in tqdm(ijson.items(f, 'item')):
                flow = json_flow['direction_length']
                # shift the negative length to positive length
                for i, length in enumerate(flow):
                    length += offset
                    length = length // 10
                    flow[i] = length
                # truncate or pad flow
                flow = np.array(flow)
                if len(flow) > truncate_num:
                    flow = flow[:truncate_num]
                elif len(flow) < truncate_num:
                    flow = np.pad(flow, ((0, truncate_num - flow.shape[0])), 'constant')
                # flow to tensor
                flow = flow.astype(np.int64)
                flow = torch.from_numpy(flow)
                
                fname = json_flow['source_file']
                site = json_flow['label']
                self.trainset.append((site, flow, fname))
    
    def __getitem__(self, index):
        return self.trainset[index]
    
    def __len__(self):
        return len(self.trainset)

